#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简化版 Elasticsearch 保险文档搜索
适用于快速测试和演示
"""

import os
import json
import logging
from datetime import datetime
from pathlib import Path

# 导入颜色高亮工具
try:
    from color_highlight_utils import process_highlight_fragment, print_colored, highlight_text
except ImportError:
    # 如果导入失败，提供简单的替代函数
    def process_highlight_fragment(fragment, style='red_bold'):
        return fragment.replace('<HIGHLIGHT>', '[').replace('</HIGHLIGHT>', ']')
    
    def print_colored(text, color='white', style=None):
        print(text)
    
    def highlight_text(text, style='red_bold'):
        return f"[{text}]"

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def check_dependencies():
    """检查依赖库"""
    missing_deps = []
    
    try:
        import elasticsearch
    except ImportError:
        missing_deps.append("elasticsearch")
    
    try:
        import PyPDF2
    except ImportError:
        missing_deps.append("PyPDF2")
    
    if missing_deps:
        print("❌ 缺少以下依赖库:")
        for dep in missing_deps:
            print(f"   pip install {dep}")
        return False
    
    return True


def simple_elasticsearch_demo():
    """简化的Elasticsearch演示"""
    
    if not check_dependencies():
        print("\n📦 请先安装依赖库，然后重新运行程序")
        return
    
    from elasticsearch import Elasticsearch
    
    print("🚀 简化版 Elasticsearch 保险文档搜索演示")
    print("="*60)
    
    # 1. 连接到Elasticsearch
    print("\n📡 步骤1: 连接到Elasticsearch")
    try:
        es = Elasticsearch(
            "https://localhost:9200",
             request_timeout=30,
            # 如果需要认证，取消下面的注释
            basic_auth=("elastic", "ncUg27UP5aMSzSElVXAX"),
            verify_certs=False,
        )
        
        if es.ping():
            print("✅ 成功连接到Elasticsearch")
            cluster_info = es.info()
            print(f"   集群: {cluster_info['cluster_name']}")
            print(f"   版本: {cluster_info['version']['number']}")
        else:
            print("❌ 无法连接到Elasticsearch")
            print("   请确保Elasticsearch服务正在运行")
            return
            
    except Exception as e:
        print(f"❌ 连接失败: {str(e)}")
        print("   请检查Elasticsearch是否已启动")
        return
    
    # 2. 创建索引
    print("\n📋 步骤2: 创建索引")
    index_name = "insurance_docs_simple"
    
    try:
        # 删除已存在的索引
        if es.indices.exists(index=index_name):
            es.indices.delete(index=index_name)
            print(f"   删除旧索引: {index_name}")
        
        # 创建新索引
        index_mapping = {
            "mappings": {
                "properties": {
                    "title": {"type": "text"},
                    "content": {"type": "text"},
                    "file_name": {"type": "keyword"},
                    "insurance_type": {"type": "keyword"},
                    "created_time": {"type": "date"}
                }
            }
        }
        
        es.indices.create(index=index_name, body=index_mapping)
        print(f"✅ 成功创建索引: {index_name}")
        
    except Exception as e:
        print(f"❌ 创建索引失败: {str(e)}")
        return
    
    # 3. 读取并索引文档
    print("\n📚 步骤3: 索引文档")
    docs_folder = "docs"
    
    if not os.path.exists(docs_folder):
        print(f"❌ 文档文件夹不存在: {docs_folder}")
        return
    
    indexed_count = 0
    
    for file_name in os.listdir(docs_folder):
        file_path = os.path.join(docs_folder, file_name)
        
        if not os.path.isfile(file_path):
            continue
        
        # 只处理txt文件（简化版）
        if not file_name.endswith('.txt'):
            continue
        
        try:
            # 读取文件内容
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # 确定保险类型
            if '雇主' in file_name:
                insurance_type = '雇主责任险'
            elif '意外' in file_name:
                insurance_type = '意外伤害保险'
            elif '责任' in file_name:
                insurance_type = '责任保险'
            else:
                insurance_type = '综合保险'
            
            # 创建文档
            doc = {
                "title": file_name.replace('.txt', ''),
                "content": content,
                "file_name": file_name,
                "insurance_type": insurance_type,
                "created_time": datetime.now().isoformat()
            }
            
            # 索引文档
            es.index(index=index_name, body=doc)
            indexed_count += 1
            print(f"   ✅ 已索引: {file_name}")
            
        except Exception as e:
            print(f"   ❌ 索引失败 {file_name}: {str(e)}")
    
    # 刷新索引
    es.indices.refresh(index=index_name)
    print(f"📊 总共索引了 {indexed_count} 个文档")
    
    # 4. 执行搜索
    print("\n🔍 步骤4: 执行搜索")
    search_query = "工伤保险和雇主险有什么区别？"
    
    try:
        search_body = {
            "query": {
                "multi_match": {
                    "query": search_query,
                    "fields": ["title^2", "content"],
                    "fuzziness": "AUTO"
                }
            },
            "highlight": {
                "fields": {
                    "content": {
                        "fragment_size": 100,
                        "number_of_fragments": 2,
                        "pre_tags": ["<HIGHLIGHT>"],
                        "post_tags": ["</HIGHLIGHT>"]
                    }
                }
            },
            "size": 5
        }
        
        response = es.search(index=index_name, body=search_body)
        
        # 5. 显示搜索结果
        print("\n📊 步骤5: 显示搜索结果")
        print("="*60)
        print(f"🔍 搜索查询: {search_query}")
        
        hits = response['hits']
        total_hits = hits['total']['value']
        
        print(f"📊 找到 {total_hits} 个相关文档")
        print("="*60)
        
        for i, hit in enumerate(hits['hits'], 1):
            source = hit['_source']
            score = hit['_score']
            
            print(f"\n📄 结果 {i} (相关度: {score:.4f})")
            print(f"标题: {source['title']}")
            print(f"文件: {source['file_name']}")
            print(f"保险类型: {source['insurance_type']}")
            
            # 显示高亮片段
            if 'highlight' in hit and 'content' in hit['highlight']:
                print("📝 相关内容:")
                for fragment in hit['highlight']['content']:
                    # 使用颜色高亮工具处理片段
                    highlighted_fragment = process_highlight_fragment(fragment, 'red_bold')
                    print(f"   {highlighted_fragment}")
            
            print("-" * 40)
        
        print(f"\n✅ 搜索完成")
        
    except Exception as e:
        print(f"❌ 搜索失败: {str(e)}")


def main():
    """主函数"""
    try:
        simple_elasticsearch_demo()
    except KeyboardInterrupt:
        print("\n👋 程序被用户中断")
    except Exception as e:
        print(f"❌ 程序执行出错: {str(e)}")


if __name__ == "__main__":
    main()